#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
task_runner.py
批处理“混元AI播客参考测试集”风格的表格（xlsx/csv），按行生成播客音频。

支持的列名（会自动模糊匹配）：
- 是否强化时效性: 可忽略
- 输入指示L0: 主题query / url / 其它
- 输入: 主题文字 或 URL
- 解析后文本: 优先使用这列作为资料注入（context）
- 输出: 可忽略

调用项目内现有组件：
- LLMClient / Role  (生成剧本)
- TTSClient / Speaker (Piper 合成)
"""

from __future__ import annotations
import argparse
import os
import re
import sys
import textwrap
from typing import List, Tuple, Optional

import pandas as pd

# 可选依赖（网页抓取用），失败时不会中断
try:
    import requests
    from bs4 import BeautifulSoup
except Exception:
    requests = None
    BeautifulSoup = None

from llm_client import LLMClient, Role
from tts_client import TTSClient, Speaker
import yaml

# 可选 RAG
try:
    from rag import retrieve as rag_retrieve
except Exception:
    rag_retrieve = None


# ---------------- 工具 ----------------
def _slug(s: str) -> str:
    s = re.sub(r"[\\/:*?\"<>|]+", "_", s.strip())
    s = re.sub(r"\s+", "_", s)
    return s[:64] if len(s) > 64 else s


def _col(df: pd.DataFrame, variants: List[str]) -> Optional[str]:
    """在 DataFrame 里找一个匹配的列名（宽松匹配）。"""
    cols = list(df.columns)
    low = {c.lower(): c for c in cols}
    for v in variants:
        if v.lower() in low:
            return low[v.lower()]
    # 次优：包含匹配
    for c in cols:
        lc = c.lower()
        if any(v.lower() in lc for v in variants):
            return c
    return None


def _read_sheet(path: str) -> pd.DataFrame:
    ext = os.path.splitext(path)[1].lower()
    if ext in (".xlsx", ".xls"):
        return pd.read_excel(path)
    elif ext == ".csv":
        return pd.read_csv(path)
    else:
        raise ValueError(f"不支持的表格格式：{ext}")


def _fetch_url_text(url: str, timeout: float = 12.0) -> Optional[str]:
    if not requests or not BeautifulSoup:
        return None
    try:
        r = requests.get(url, timeout=timeout, headers={"User-Agent": "Mozilla/5.0"})
        if r.status_code != 200:
            return None
        soup = BeautifulSoup(r.text, "lxml")
        # 简单抽取正文
        for tag in soup(["script", "style", "noscript"]):
            tag.extract()
        text = soup.get_text("\n")
        text = re.sub(r"\n{2,}", "\n", text).strip()
        return text[:8000]  # 控长
    except Exception:
        return None


def _load_roles_yaml(path: str) -> Tuple[List[Role], List[Speaker]]:
    data = yaml.safe_load(open(path, "r", encoding="utf-8"))
    roles_cfg = data.get("roles", [])
    roles: List[Role] = []
    speakers: List[Speaker] = []
    for rc in roles_cfg:
        name = rc["name"]
        style = rc.get("style", "")
        backend = rc.get("backend", "piper").lower()
        if backend != "piper":
            raise ValueError(f"仅支持 piper，发现 backend={backend}（角色={name}）")

        roles.append(Role(name=name, style=style))

        speakers.append(
            Speaker(
                name=name,
                backend="piper",
                model_path=rc.get("model_path"),
                config_path=rc.get("config_path"),
                sample_rate=int(rc.get("sample_rate", 22050)),
                sentence_silence=float(rc.get("sentence_silence", 0.5)),
                tts_params=rc.get("tts", {}) or {},
                cmd_template=rc.get("cmd"),
            )
        )
    return roles, speakers


# ---------------- 主流程 ----------------
def run_sheet(
    sheet_path: str,
    roles_yaml: str,
    out_root: str,
    minutes: int = 3,
    mode: str = "normal",
    rounds: Optional[int] = None,
    rows: Optional[str] = None,
):
    df = _read_sheet(sheet_path)

    col_l0 = _col(df, ["输入指示L0", "指示", "类型"])
    col_in = _col(df, ["输入", "内容", "主题", "url"])
    col_parsed = _col(df, ["解析后文本", "解析文本", "摘要", "资料"])
    col_id = _col(df, ["序号", "id", "index"])

    if not col_in:
        raise ValueError("表格里没有找到“输入/主题/url”等列。请检查列名。")

    roles, speakers = _load_roles_yaml(roles_yaml)
    client = LLMClient()

    # 选择行
    if rows:
        selected = set()
        for part in rows.split(","):
            part = part.strip()
            if "-" in part:
                a, b = part.split("-", 1)
                selected.update(range(int(a), int(b) + 1))
            else:
                selected.add(int(part))
        df = df.iloc[[i - 1 for i in sorted(selected) if 1 <= i <= len(df)]]

    for ridx, row in df.iterrows():
        l0 = str(row.get(col_l0, "") or "").strip()
        raw_in = str(row.get(col_in, "") or "").strip()
        parsed = str(row.get(col_parsed, "") or "").strip()
        sid = str(row.get(col_id, "") or (ridx + 1))

        if not raw_in and not parsed:
            print(f"[skip] 第{sid}行：没有输入也没有解析后文本。")
            continue

        # 主题 & 注入资料
        topic = ""
        context = None

        if parsed:
            # 已有“解析后文本” → 作为 context，topic 取第一行标题或输入
            context = parsed
            topic = (raw_in or parsed.splitlines()[0])[:60]
        else:
            if l0.lower().startswith("url") or raw_in.startswith("http"):
                topic = f"解析与讨论：{raw_in}"
                webtxt = _fetch_url_text(raw_in)
                context = webtxt or None
            else:
                topic = raw_in

        if mode == "deep" and rag_retrieve is not None:
            # 叠加 RAG（如果 context 已有会一起注入）
            rag_ctx = rag_retrieve(topic, top_k=3)
            context = (context or "") + "\n\n" + (rag_ctx or "")

        # 生成脚本
        script = client.generate_dialogue(
            topic=topic,
            roles=roles,
            minutes=minutes,
            rounds=rounds,
            mode=mode,
            context=context,
        )

        # 输出目录
        run_dir = os.path.join(
            out_root, f"{int(sid) if sid.isdigit() else sid:>03}_{_slug(topic)}"
        )
        os.makedirs(run_dir, exist_ok=True)

        # 存脚本
        header = textwrap.dedent(
            f"""\
            # 行ID：{sid}
            # 主题：{topic}
            # 模式：{mode}
            # 角色：{", ".join(r.name for r in roles)}
            """
        )
        script_path = os.path.join(run_dir, "script.txt")
        with open(script_path, "w", encoding="utf-8") as f:
            f.write(header + "\n" + script)

        # 合成音频
        tts = TTSClient(speakers=speakers, out_dir=os.path.join(run_dir, "tts"))
        segs = tts.synthesize_dialogue(header + script, os.path.join(run_dir, "tts"))

        # 合并与转码
        from audio_utils import concat_wavs_ffmpeg, to_mp3_ffmpeg

        final_wav = concat_wavs_ffmpeg(segs, os.path.join(run_dir, "podcast.wav"))
        final_mp3 = to_mp3_ffmpeg(final_wav, os.path.join(run_dir, "podcast.mp3"))

        print(f"[ok] 第{sid}行完成：脚本={script_path} 音频={final_wav} / {final_mp3}")


def main():
    ap = argparse.ArgumentParser(description="表格批处理生成播客（Piper）")
    ap.add_argument("--sheet", required=True, help="xlsx 或 csv 文件路径")
    ap.add_argument("--roles", default="configs/roles_piper.yaml", help="角色配置YAML")
    ap.add_argument("--outdir", default="output_batch", help="输出根目录")
    ap.add_argument("--minutes", type=int, default=3)
    ap.add_argument("--mode", default="normal", choices=["normal", "deep"])
    ap.add_argument("--rounds", type=int, default=None)
    ap.add_argument("--rows", default=None, help="行号选择，如 '3-5,7,10'（1-based）")
    args = ap.parse_args()

    run_sheet(
        sheet_path=args.sheet,
        roles_yaml=args.roles,
        out_root=args.outdir,
        minutes=args.minutes,
        mode=args.mode,
        rounds=args.rounds,
        rows=args.rows,
    )


if __name__ == "__main__":
    main()
